# Module containing CUDA code.  At this point functions for rotating,
# forward projecting and back projecting.
#


from pycuda.compiler import SourceModule

projection_module = SourceModule("""

__global__ void ForwardProject1(float * V, float * P, float * R, float radius,
		int N, int M, int nstep)
	{
	float radius2 = 0.5f*float(N)*radius;
	radius2 = radius2*radius2;

	int index = N*blockIdx.x + threadIdx.x;

	float centre = 0.5f*float(N);
	float shift = 1.0f/float(M);
	float3 crd;
	float stride = 1.0f/float(M);

	for (int i=0; i<M; i++) {
		for (int j=0; j<M; j++) {
			crd.x = blockIdx.x + i*shift + 0.5*shift - centre; 
			crd.y = threadIdx.x + j*shift + 0.5*shift - centre; 
			crd.z = radius2 - crd.x*crd.x - crd.y*crd.y;

			// is there a more effecient way to generate a weight
			float weight = ceilf(crd.z);
			weight = min(1.0f,weight);
			weight = max(0.0f,weight);

			// if x,y are outside the radius, point is thrown to the centre
			crd.x *= weight;
			crd.y *= weight;
			crd.z = sqrtf(weight*crd.z);
			
			float3 stride;
			stride.x = R[6]*2.0f*crd.z/float(nstep);
			stride.y = R[7]*2.0f*crd.z/float(nstep);
			stride.z = R[8]*2.0f*crd.z/float(nstep);

			float3 start;
			start.x = R[0]*crd.x + R[3]*crd.y - R[6]*crd.z;
			start.y = R[1]*crd.x + R[4]*crd.y - R[7]*crd.z;
			start.z = R[2]*crd.x + R[5]*crd.y - R[8]*crd.z;

			int3 vcrd;
			for (int k=0; k<nstep; k++) {
				vcrd.x = int(start.x + centre);
				vcrd.y = int(start.y + centre);
				vcrd.z = int(start.z + centre);
				P[index] += weight*V[N*N*vcrd.x + N*vcrd.y + vcrd.z];

				start.x += stride.x;
				start.y += stride.y;
				start.z += stride.z;
				}
			}
		}

	P[index] /= float(M*M);
	}

__global__ void ForwardProjectMany(float * V, float * P, float * R,
		float radius, int N, int M, int nstep)
	{
	float radius2 = 0.5f*float(N)*radius;
	radius2 = radius2*radius2;
	
	int pindex = N*N*blockIdx.x + N*blockIdx.y + threadIdx.x;
	float * rot = &R[9*blockIdx.x];

	float centre = 0.5f*float(N);
	float shift = 1.0f/float(M);
	float3 crd;

	for (int i=0; i<M; i++) {
		for (int j=0; j<M; j++) {
			crd.x = blockIdx.y + i*shift + 0.5*shift - centre; 
			crd.y = threadIdx.x + j*shift + 0.5*shift - centre; 
			crd.z = radius2 - crd.x*crd.x - crd.y*crd.y;

			// is there a more effecient way to generate a weight
			float weight = ceilf(crd.z);
			weight = min(1.0f,weight);
			weight = max(0.0f,weight);

			// if x,y are outside the radius, point is thrown to the centre
			crd.x *= weight;
			crd.y *= weight;
			crd.z = sqrtf(weight*crd.z);
			
			float3 stride;
			stride.x = rot[6]*2.0f*crd.z/float(nstep);
			stride.y = rot[7]*2.0f*crd.z/float(nstep);
			stride.z = rot[8]*2.0f*crd.z/float(nstep);

			float3 start;
			start.x = rot[0]*crd.x + rot[3]*crd.y - rot[6]*crd.z;
			start.y = rot[1]*crd.x + rot[4]*crd.y - rot[7]*crd.z;
			start.z = rot[2]*crd.x + rot[5]*crd.y - rot[8]*crd.z;

			int3 vcrd;
			float raysum = 0.0f;
			for (int k=0; k<nstep; k++) {
				vcrd.x = int(start.x + centre);
				vcrd.y = int(start.y + centre);
				vcrd.z = int(start.z + centre);
				raysum += V[N*N*vcrd.x + N*vcrd.y + vcrd.z];
//				P[pindex] += weight*V[N*N*vcrd.x + N*vcrd.y + vcrd.z];

				start.x += stride.x;
				start.y += stride.y;
				start.z += stride.z;
				}
			raysum *= 2.0*crd.z/float(nstep);
			P[pindex] += weight*raysum;
			}
		}

// need to fix adding of the rays
//				P[pindex] *= 2.0*crd.z/float(nstep);
	P[pindex] /= float(M*M);
	}


// Back project a single image over a volume specified by a radius
// of interpolation points
// V -> volume to project
// P -> projection plane
// R -> rotation matrix in row major order
// N -> size of the projection plane/volume, and volume in pixels/voxels
// M -> MxM division in a pixel of the projection plane
// radius -> radius of the volume

__global__ void BackProject1(float * V, float * P, float * R,
		int N, int M, float radius)
	{
	// radius squared given in pixels/voxels
	float radius2 = 0.5*float(N)*radius;
	radius2 = radius2*radius2;

	// this is which voxel, blocks of threads go in the direction of z
	int vidx = N*N*blockIdx.x + N*blockIdx.y + threadIdx.x;

	// a voxel is divided into MxMxM partial voxels.  it takes M strides
	// to cross the edge of one voxel.  
	float stride = 1.0f/float(M);

	// for finding the centre of a partial voxel.
	float shift = 0.5*stride;

	// the centre of the image
	float centre = 0.5*float(N);

	// coordinate of the centre of the partial voxel to be ready for rotation
	float3 crd;


	for (int i=0; i<M; i++) {
		crd.x = blockIdx.x - centre + float(i)*stride + shift;
		for (int j=0; j<M; j++) {
			crd.y = blockIdx.y - centre + float(j)*stride + shift;
			for (int k=0; k<M; k++) {
				crd.z = threadIdx.x - centre + float(k)*stride + shift;

				// rotated coordinate of the partial voxel
				float3 xyz;
				xyz.x = R[0]*crd.x + R[1]*crd.y + R[2]*crd.z;
				xyz.y = R[3]*crd.x + R[4]*crd.y + R[5]*crd.z;
				xyz.z = R[6]*crd.x + R[7]*crd.y + R[8]*crd.z;

				// trick to give the voxel a weight of zero or one 
				// depending on whether it is in or out of the radius.
				float vradius2 = xyz.x*xyz.x + xyz.y*xyz.y + xyz.z*xyz.z;
				float weight = floorf(vradius2/radius2);
				weight = 1.0f - min(weight,1.0f);

				// trick to locate the sampled value from within the
				// projection array.
				int row = max(0,min(N-1,int(xyz.x + centre)));
				int col = max(0,min(N-1,int(xyz.y + centre)));

				// the length of the projection path is at least 1.0f/M
				float raylen = radius2 - xyz.x*xyz.x - xyz.y*xyz.y;
				raylen = sqrt(max(stride*stride,2.0*raylen));

				// grab a sample
				V[vidx] += weight*P[N*row + col]/raylen;
				}
			}
		}

	// divide by the number of partial volumes
	V[vidx] *= stride*stride*stride;
	}


__global__ void BackProjectMany(float * V, float * P, float * R,
		int nproj, int N, int M, float radius)
	{
	// radius squared given in pixels/voxels
	float radius2 = 0.5*float(N)*radius;
	radius2 = radius2*radius2;

	// this is which voxel, blocks of threads go in the direction of z
	int vidx = N*N*blockIdx.x + N*blockIdx.y + threadIdx.x;

	// a voxel is divided into MxMxM partial voxels.  it takes M strides
	// to cross the edge of one voxel.  
	float stride = 1.0f/float(M);

	// for finding the centre of a partial voxel.
	float shift = 0.5*stride;

	// the centre of the image
	float centre = 0.5*float(N);

	// coordinate of the centre of the partial voxel to be ready for rotation
	float3 crd;

	float add_val = 0.0f;

	for (int i=0; i<M; i++) {
		crd.x = blockIdx.x - centre + float(i)*stride + shift;
		for (int j=0; j<M; j++) {
			crd.y = blockIdx.y - centre + float(j)*stride + shift;
			for (int k=0; k<M; k++) {
				crd.z = threadIdx.x - centre + float(k)*stride + shift;

				// trick to give the voxel a weight of zero or one 
				// depending on whether it is in or out of the radius.
				float vradius2 = crd.x*crd.x + crd.y*crd.y + crd.z*crd.z;
				float weight = floorf(vradius2/radius2);
				weight = 1.0f - min(weight,1.0f);

				for (int t=0; t<nproj; t++) {
					// rotated coordinate of the partial voxel
					float * rot = &R[9*t];
					float3 xyz;
					xyz.x = rot[0]*crd.x + rot[1]*crd.y + rot[2]*crd.z;
					xyz.y = rot[3]*crd.x + rot[4]*crd.y + rot[5]*crd.z;
					xyz.z = rot[6]*crd.x + rot[7]*crd.y + rot[8]*crd.z;

					// trick to locate the sampled value from within the
					// projection array.
					int row = max(0,min(N-1,int(xyz.x + centre)));
					int col = max(0,min(N-1,int(xyz.y + centre)));

					// the length of the projection path is at least 1.0f/M
					float raylen = radius2 - xyz.x*xyz.x - xyz.y*xyz.y;
					raylen = sqrt(max(stride*stride,4.0*raylen));

					// grab a sample
					add_val += weight*P[N*N*t + N*row + col]/raylen;
					}
				}
			}
		}

	// divide by the number of partial volumes
	add_val /= nproj;
	add_val *= stride*stride*stride;
	V[vidx] += add_val;
	}
	""")

ForwardProject1 = projection_module.get_function("ForwardProject1")
ForwardProjectMany = projection_module.get_function("ForwardProjectMany")
BackProject1 = projection_module.get_function("BackProject1")
BackProjectMany = projection_module.get_function("BackProjectMany")

